"""
水平分表实现模块
支持按照ID取模实现数据分片存储
提供统一的查询和保存接口
"""
from django.db import models, connections
from django.db.models.sql import Query
from .snowflake import generate_id


class ShardedManager(models.Manager):
    """分表模型管理器"""
    
    def get_table_name(self, table_id):
        """获取实际表名"""
        return f"{self.model._meta.db_table}_{table_id}"
    
    def get_table_id(self, obj_id):
        """
        通过对象ID确定表ID (对3取模)
        
        Args:
            obj_id: 对象ID
        
        Returns:
            int: 表ID (0-2)
        """
        return obj_id % 3
        
    def get_queryset_for_table(self, table_id):
        """
        获取特定表的查询集
        
        Args:
            table_id: 表ID
            
        Returns:
            QuerySet: 对应表的查询集
        """
        queryset = super().get_queryset()
        queryset.model._meta.db_table = self.get_table_name(table_id)
        return queryset
    
    def save_to_table(self, obj, table_id):
        """
        将对象保存到指定表
        
        Args:
            obj: 要保存的对象
            table_id: 表ID
        """
        obj._meta.db_table = self.get_table_name(table_id)
        obj.save(using=obj._state.db or 'default')
    
    def create(self, **kwargs):
        """
        创建新对象并保存
        
        Args:
            **kwargs: 对象属性
        
        Returns:
            Object: 创建的对象实例
        """
        # 生成雪花ID
        if 'id' not in kwargs:
            kwargs['id'] = generate_id()
            
        # 确定分表ID
        table_id = self.get_table_id(kwargs['id'])
        
        # 创建对象但不保存
        obj = self.model(**kwargs)
        
        # 保存到指定表
        self.save_to_table(obj, table_id)
        return obj
    
    def get(self, **kwargs):
        """
        从分表中获取对象
        
        Args:
            **kwargs: 查询条件
        
        Returns:
            Object: 查找到的对象
        
        Raises:
            DoesNotExist: 对象不存在
            MultipleObjectsReturned: 多个对象满足条件
        """
        # 如果有id，直接根据id确定表
        if 'id' in kwargs:
            table_id = self.get_table_id(kwargs['id'])
            queryset = self.get_queryset_for_table(table_id)
            return queryset.get(**kwargs)
        
        # 否则需要查询所有分表
        for i in range(3):
            queryset = self.get_queryset_for_table(i)
            try:
                return queryset.get(**kwargs)
            except self.model.DoesNotExist:
                continue
        
        # 所有表都没找到，抛出异常
        raise self.model.DoesNotExist(
            f"{self.model.__name__} matching query does not exist."
        )
    
    def filter(self, **kwargs):
        """
        从分表中筛选对象
        
        Args:
            **kwargs: 筛选条件
        
        Returns:
            List: 匹配的对象列表
        """
        # 如果有id，直接根据id确定表
        if 'id' in kwargs:
            table_id = self.get_table_id(kwargs['id'])
            queryset = self.get_queryset_for_table(table_id)
            return queryset.filter(**kwargs)
        
        # 否则需要合并所有分表的结果
        results = []
        for i in range(3):
            queryset = self.get_queryset_for_table(i)
            results.extend(list(queryset.filter(**kwargs)))
        
        return results
    
    def all(self):
        """
        获取所有分表的对象
        
        Returns:
            List: 所有表的对象列表
        """
        results = []
        for i in range(3):
            queryset = self.get_queryset_for_table(i)
            results.extend(list(queryset.all()))
        
        return results


class ShardedModel(models.Model):
    """分表模型基类"""
    
    objects = ShardedManager()
    
    class Meta:
        abstract = True
    
    def save(self, *args, **kwargs):
        """重写保存方法，实现分表存储"""
        # 如果是新对象且没有ID，生成雪花ID
        if not self.id:
            self.id = generate_id()
            
        # 确定分表ID
        table_id = self.__class__.objects.get_table_id(self.id)
        
        # 设置实际表名
        self._meta.db_table = self.__class__.objects.get_table_name(table_id)
        
        # 调用原始保存方法
        super().save(*args, **kwargs)


class ShardedOrder(ShardedModel):
    """订单分表基类"""
    
    class Meta:
        abstract = True
    
    @staticmethod
    def create_tables():
        """创建分表，数据库迁移时使用"""
        from django.db import connection
        
        # 获取基本建表SQL
        base_sql = "CREATE TABLE IF NOT EXISTS {table_name} LIKE background_order;"
        
        # 创建三个分表
        with connection.cursor() as cursor:
            for i in range(3):
                table_name = f"background_order_{i}"
                sql = base_sql.format(table_name=table_name)
                cursor.execute(sql) 